import os
import paramiko
import getpass
import hashlib
import pickle
from .authenticators import AuthenticationError


class HostIdentityError(Exception):
    pass
class UnknownHostIdentity(HostIdentityError):
    pass
class HostIdentityChanged(HostIdentityError):
    pass


class ParamikoClient(object):
    KNOWN_HOSTS_FILES = ['~/.ssh/known_hosts', '~/ssh/known_hosts']
    ALLOW_UNKNOWN_HOSTS = True
    
    def __init__(self, sock, username, hostname = None):
        self._username = username
        self._hostname = hostname if hostname else sock.getpeername()[0]
        self._transport = paramiko.Transport(sock)
        self._chan = None
        
        self._transport.start_client()
        self._hostkey = self._transport.get_remote_server_key()
        self._known_hosts = self.get_known_hosts(self.KNOWN_HOSTS_FILES)
    
    @classmethod
    def connect(cls, sock, username, password = None, key = None, rsa_file = None, dsa_file = None):
        inst = cls(sock, username)
        inst.validate_host_identity(inst.ALLOW_UNKNOWN_HOSTS)
        if password:
            inst.authenticate_password(password)
        elif key:
            self.authenticate_key(key)
        elif rsa_file:
            self.authenticate_key(rsa_file)
        elif dsa_file:
            self.authenticate_key(dsa_file)
        else:
            raise ValueError("no authentication methods available")
        return inst.getChannel()
    
    def __del__(self):
        self.close()
    
    def close(self):
        if self._chan:
            self._chan.close()
            self._chan = None
        if self._transport:
            self._transport.close()
            self._transport = None
    
    @classmethod
    def get_known_hosts(cls, files):
        for fn in files:
            try:
                keys = paramiko.util.load_host_keys(os.path.expanduser(fn))
                return keys
            except IOError:
                pass
        return {}
    
    def validate_host_identity(self, allow_unknown_hosts = False, allow_changed_hosts = False):
        if self._hostname not in self._known_hosts:
            if not allow_unknown_hosts:
                raise UnknownHostIdentity()
        elif key.get_name() not in self._known_hosts[self._hostname]:
            if not allow_unknown_hosts:
                raise UnknownHostIdentity()
        elif self._known_hosts[self._hostname] != self._hostkey:
            if not allow_changed_hosts:
                raise HostIdentityChanged()
        else:
            pass # all ok

    def authenticate_with_agent(self):
        agent = paramiko.Agent()
        keys = agent.get_keys()
        if not keys:
            raise paramiko.SSHException("agent found no keys")
        for key in keys:
            try:
                self._transport.auth_publickey(self._username, key)
                assert self._transport.is_authenticated()
                return
            except paramiko.SSHException, ex:
                pass
        raise ex

    def authenticate_interactive(self):
        pw = getpass.getpass('Password for %s@%s: ' % (self._username, self._hostname))
        self._transport.auth_password(self._username, pw)
        assert self._transport.is_authenticated()

    def authenticate_password(self, passowrd):
        self._transport.auth_password(self._username, passowrd)
        assert self._transport.is_authenticated()
    
    def authenticate_key(self, key):
        self._transport.auth_publickey(self._username, key)
        assert self._transport.is_authenticated()
    
    def authenticate_rsa(self, filename, password = None):
        try:
            key = paramiko.RSAKey.from_private_key_file(filename)
        except paramiko.PasswordRequiredException:
            if password is None:
                password = getpass.getpass('RSA key password (%s): ' % (filename,))
            key = paramiko.RSAKey.from_private_key_file(filename, password)
        self.authenticate_key(key)

    def authenticate_dsa(self, filename, password = None):
        try:
            key = paramiko.DSSKey.from_private_key_file(filename)
        except paramiko.PasswordRequiredException:
            if password is None:
                password = getpass.getpass('DSA key password (%s): ' % (filename,))
            key = paramiko.DSSKey.from_private_key_file(filename, password)
        self.authenticate_key(key)
    
    def getChannel(self):
        if self._chan is None:
            self._chan = self._transport.open_channel("raw")
        return self._chan


class ParamikoAuthenticator(paramiko.ServerInterface):
    def __init__(self, users, host_key = None, filename = None):
        self.users = users
        self.host_key = paramiko.RSAKey(filename = "/usr/share/doc/python-paramiko/examples/test_rsa.key")
        self.filename = filename

    @classmethod
    def from_file(cls, filename):
        if os.path.exists(filename):
            f = open(filename, "rb")
            users = pickle.load(f)
            f.close()
        else:
            users = {}
        return cls(users, filename)

    @classmethod
    def _get_hashed_password(cls, username, password):
        return hashlib.sha1(username + password).digest()
    
    def set_user_password(self, username, password):
        if username not in self.users:
            self.users[username] = [None, None]
        self.users[username][0] = self._get_hashed_password(username, password)

    def set_user_private_key(self, username, key = None, key_file = None):
        if not key:
            if isinstance(key_file, basestring):
                key_file = open(key_file, "rb")
            key = paramiko.RSAKey.from_private_key(key_file)
        if username not in self.users:
            self.users[username] = [None, None]
        self.users[username][1] = key

    def del_user(self, username):
        del self.users[username]

    def list_users(self):
        return self.users.keys()

    def sync(self):
        f = open(self.filename, "wb")
        data = pickle.dump(self.users, f)
        f.close()

    @classmethod
    def _get_credentials(cls, transport):
        return transport.get_username(), transport.get_server_key()

    def __call__(self, sock):
        #try:
        t = paramiko.Transport(sock)
        t.add_server_key(self.host_key)
        t.start_server(server = self)
        chan = t.accept()
        #except Exception, ex:
        #    raise AuthenticationError(str(ex))
        return chan, self._get_credentials(t)

    #
    # paramiko.ServerInterface
    #
    def check_channel_request(self, kind, chanid):
        if kind == "raw":
            return paramiko.OPEN_SUCCEEDED
        return paramiko.OPEN_FAILED_ADMINISTRATIVELY_PROHIBITED

    def get_allowed_auths(self, username):
        return "password,publickey"

    def check_auth_password(self, username, password):
        pw, pk = self.users.get(username, (None, None))
        if pw == self._get_hashed_password(username, password):
            return paramiko.AUTH_SUCCESSFUL
        else:
            return paramiko.AUTH_FAILED

    def check_auth_publickey(self, username, key):
        pw, pk = self.users.get(username, (None, None))
        if pk == key:
            return paramiko.AUTH_SUCCESSFUL
        else:
            return paramiko.AUTH_FAILED




